package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
)

type RenderArgs struct {
	HtmlFilename string
	PDFFilename  string
	PageWidth    int
	PageHeight   int
	MarginTop    int
	MarginBottom int
	MarginLeft   int
	MarginRight  int
}

//
// channel for communication
//
// ch <- wkhtmltopdf arguments
// <- ch DONE/ERROR
//
// whenever ERROR occurs, this routine would be dead
//
func ForkRun(ch chan string) {
	logger := log.New(os.Stdout, " PDF: ", log.LstdFlags)
	cmd := exec.Command("xvfb-run", "wkhtmltopdf", "--minimum-font-size", "1", "--read-args-from-stdin")
	defer cmd.Wait()
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		logger.Printf("Error: %v", err)
		ch <- "ERROR"
		return
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		logger.Printf("Error: %v\n", err)
		ch <- "ERROR"
		return
	}
	cmd.Start()

	breader := bufio.NewReader(stdout)
	for {
		args := <-ch
		logger.Printf("Received: %v", args)
		_, err := io.WriteString(stdin, args)
		if err != nil {
			// broken pipe?
			logger.Printf("Error: %v", err)
			ch <- "ERROR"
			break
		}
		success := true
		for success {
			b, err := breader.ReadString('\n')
			logger.Printf("STDOUT: %s", b)
			if err != nil {
				logger.Printf("Error: %v", err)
				success = false
				ch <- "ERROR"
				break
			}
			if strings.Contains(b, "Done") {
				ch <- "DONE"
				break
			} else if strings.Contains(b, "<Error>") {
				success = false
				cmd.Process.Kill()
				ch <- "FATAL"
				break
			}
		}
		if !success {
			break
		}
		logger.Printf("Done: %v", args)
	}
}

func NewRenderArgs(htmlFilename, pdfFilename string,
	width, height, top, bottom, left, right int) RenderArgs {

	return RenderArgs{HtmlFilename: htmlFilename,
		PDFFilename:  pdfFilename,
		PageWidth:    width,
		PageHeight:   height,
		MarginTop:    top,
		MarginBottom: bottom,
		MarginLeft:   left,
		MarginRight:  right,
	}
}

func MakeCommand(args RenderArgs) string {
	cmd := fmt.Sprintf("--encoding utf-8 --page-width %dmm --page-height %dmm --margin-top %dmm --margin-bottom %dmm --margin-left %dmm --margin-right %dmm %s %s\n",
		args.PageWidth, args.PageHeight, args.MarginTop, args.MarginBottom, args.MarginLeft, args.MarginRight,
		args.HtmlFilename, args.PDFFilename)
	return cmd
}

func NewWorker() chan string {
	ch := make(chan string)
	go ForkRun(ch)
	return ch
}

type WorkerPool struct {
	lock   sync.Mutex
	avail  []chan string
	logger *log.Logger
}

func NewWorkerPool() *WorkerPool {
	wp := new(WorkerPool)
	wp.avail = append(wp.avail, NewWorker())
	wp.logger = log.New(os.Stdout, "POOL: ", log.LstdFlags)
	return wp
}

func (wp *WorkerPool) RequestWorker() chan string {
	wp.lock.Lock()
	defer wp.lock.Unlock()
	wp.logger.Printf("Requesting: Avail=%d", len(wp.avail))
	if len(wp.avail) > 0 {
		x, avail := wp.avail[len(wp.avail)-1], wp.avail[:len(wp.avail)-1]
		wp.avail = avail
		return x
	} else {
		return NewWorker()
	}
}

func (wp *WorkerPool) Recycle(ch chan string) {
	wp.lock.Lock()
	defer wp.lock.Unlock()
	wp.avail = append(wp.avail, ch)
	wp.logger.Printf("Recycled: Avail=%d", len(wp.avail))
}

func parseInt(s string, dft int) int {
	if s == "" {
		return dft
	}
	i, err := strconv.Atoi(s)
	if err != nil {
		return dft
	}
	return i
}

func parseHeader(req *http.Request) (int, int, int, int, int, int) {
	width := parseInt(req.Header.Get("Page-Width"), 100)
	height := parseInt(req.Header.Get("Page-Height"), 100)
	top := parseInt(req.Header.Get("Margin-Top"), 0)
	bottom := parseInt(req.Header.Get("Margin-Bottom"), 0)
	left := parseInt(req.Header.Get("Margin-Left"), 0)
	right := parseInt(req.Header.Get("Margin-Right"), 0)
	return width, height, top, bottom, left, right
}

func HtmlToPdfHandler(pool *WorkerPool, logger *log.Logger) func(http.ResponseWriter, *http.Request) {
	return func(resp http.ResponseWriter, req *http.Request) {
		logger.Printf("RequestURI: %v", req.RequestURI)

		var htmlfilename string

		switch req.Method {
		case "POST":
			html, err := ioutil.TempFile("/tmp", "html")
			if err != nil {
				resp.WriteHeader(500)
				return
			}
			logger.Printf("HTML: %v", html.Name())
			io.Copy(html, req.Body)
			htmlfilename = html.Name() + ".html"
			os.Rename(html.Name(), htmlfilename)
			break
		case "GET":
			htmlfilename = req.FormValue("url")
			if htmlfilename == "" {
				logger.Printf("GET: url not given")
				resp.WriteHeader(304)
				return
			}
			break
		default:
			logger.Printf("Method not supported: %s", req.Method)
			resp.WriteHeader(304)
			return
		}

		// PDF
		width, height, top, bottom, left, right := parseHeader(req)
		pdf, err := ioutil.TempFile("/tmp", "pdf")
		if err != nil {
			resp.WriteHeader(500)
			return
		}

		// error should be retryable
		retry := 3
		success := false
		for retry > 0 {
			ch := pool.RequestWorker()
			ch <- MakeCommand(NewRenderArgs(htmlfilename, pdf.Name(), width, height, top, bottom, left, right))
			flag := <-ch
			if flag == "DONE" {
				resp.Header().Set("Content-Type", "application/pdf") // normal header
				io.Copy(resp, pdf)
				os.Remove(htmlfilename)
				os.Remove(pdf.Name())
				pool.Recycle(ch)
				success = true
				break
			} else if flag == "FATAL" {
				logger.Printf("Fatal Error! No Retry")
				success = false
				break
			}
			retry--
			logger.Printf("Error!  Retry Left: %d", retry)
		}

		if !success {
			logger.Printf("Conversion FAILED!")
			resp.WriteHeader(500)
		}
	}
}

func main() {
	mux := http.NewServeMux()
	pool := NewWorkerPool()
	logger := log.New(os.Stdout, "HTTP: ", log.LstdFlags)
	mux.HandleFunc("/", HtmlToPdfHandler(pool, logger))
	http.ListenAndServe(":8080", mux)
}
